"""
独热向量的维度会映射到非常大的维度特征,为了解决维度诅咒,这里使用embed的方式替换独热向量,
将input的大维度(稀疏、硬编码)映射到embed的维度(就是input的特征作为行,embed作为列,
此时就是将input的高维度特征替换、映射到embed的低维度特征)

本例目标是将hello->ohlol
"""

import torch
import torch.nn as nn

# 输入维度特征数
input_size = 4
# 隐层的维度
hidden_size = 8
# 嵌入层的维度
embed_size = 10
# rnn的层数
num_layers = 2
# 一批的数
batch_size = 1

# 定义字符表,通过索引可以映射到字符
idx2char = ['h', 'e', 'l', 'o']
# 输入,hello,对应的idx2char中的索引
x_data = [0, 1, 2, 2, 3]
# 参照值,ohlol
y_data = [3, 0, 2, 3, 2]

# 转换到tensor
# 输入给到embed时要求数据类型为long,此时是作为索引,检索embed的矩阵某行数据
inputs = torch.tensor(x_data, dtype=torch.long).view(batch_size, -1)
# y_data转换
labels = torch.tensor(y_data, dtype=torch.long)


class MyEmbedRNN(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # 定义embedding层
        self.embed = nn.Embedding(input_size, embed_size)
        # 定义RNN,因为embed出来的数据格式为(batchSize,seqSize,embedSize),所以用batchFirst模式
        self.rnn = nn.RNN(embed_size, hidden_size, num_layers, batch_first=True)
        # 定义fc全链接网络,输出到目标维度
        self.linear = nn.Linear(hidden_size, input_size)

    def forward(self, input):
        # 经过嵌入层变化,此时输出为(batchSize,seqSize,embedSize)
        embed = self.embed(input)
        # 初始化rnn的h0输入层
        h0 = torch.zeros(num_layers, batch_size, hidden_size)
        # 进行rnn层计算
        output, _ = self.rnn(embed, h0)
        # 进行全链接线性层变化,使用交叉熵做loss计算,不需要激活层
        output = self.linear(output)
        # 转换维度(input_size即为最终特征数),用于loss计算
        return output.view(-1, input_size)


# 初始化模型
model = MyEmbedRNN()
# loss计算,使用交叉熵
criterion = nn.CrossEntropyLoss()
# 优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

for epoch in range(20):
    # 清除梯度
    optimizer.zero_grad()
    # 前馈计算
    y_pred = model(inputs)
    # 损失计算
    loss = criterion(y_pred, labels)
    # 反向传播
    loss.backward()
    # 权重跟新
    optimizer.step()

    # 当前的输出结果
    # 检索输出的最大值,用来进行映射
    _, idx = y_pred.max(dim=1)
    print(f"epoch: {epoch}, loss: {loss.item()}, {''.join([idx2char[i] for i in idx])}")
